## MIT code from https://github.com/jxnl/blog/blob/main/generate_sitemap.py

import os
import asyncio
import yaml
from typing import Generator, Tuple, Dict, Optional
from openai import AsyncOpenAI
import typer
from rich.console import Console
from rich.progress import Progress
import hashlib
from dotenv import load_dotenv

load_dotenv()

console = Console()


def traverse_docs(
    root_dir: str = "docs",
) -> Generator[Tuple[str, str, str], None, None]:
    """
    Recursively traverse the docs folder and yield the path, content, and content hash of each file.

    Args:
        root_dir (str): The root directory to start traversing from. Defaults to 'docs'.

    Yields:
        Tuple[str, str, str]: A tuple containing the relative path from 'docs', the file content, and the content hash.
    """
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith(".md"):  # Assuming we're only interested in Markdown files
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, root_dir)

                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()

                content_hash = hashlib.md5(content.encode()).hexdigest()
                yield relative_path, content, content_hash


# async def summarize_content(client: AsyncOpenAI, path: str, content: str) -> str:
#     """
#     Summarize the content of a file.

#     Args:
#         client (AsyncOpenAI): The AsyncOpenAI client.
#         path (str): The path of the file.
#         content (str): The content of the file.

#     Returns:
#         str: A summary of the content.
#     """
#     try:
#         response = await client.chat.completions.create(
#             model="gpt-4o-mini",
#             messages=[
#                 {
#                     "role": "system",
#                     "content": "You are a helpful assistant that summarizes text.",
#                 },
#                 {"role": "user", "content": content},
#                 {
#                     "role": "user",
#                     "content": "Please summarize the content in a few sentences so they can be used for SEO. Include core ideas, objectives, and important details and key points and key words",
#                 },
#             ],
#             max_tokens=4000,
#         )
#         return response.choices[0].message.content
#     except Exception as e:
#         console.print(f"[bold red]Error summarizing {path}: {str(e)}[/bold red]")
#         return ""

async def summarize_content(client: AsyncOpenAI, path: str, content: str) -> str:
    """
    Summarize the content of a file with SEO optimization for AI Engineering Academy.

    Args:
        client (AsyncOpenAI): The AsyncOpenAI client.
        path (str): The path of the file.
        content (str): The content of the file.

    Returns:
        str: An SEO-optimized summary of the content.
    """
    try:
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": """You are an expert SEO content analyzer for AI Engineering Academy, a premier educational platform focused on teaching practical AI engineering skills, machine learning implementation, and MLOps best practices. 
                    
Your task is to create SEO-optimized meta descriptions that will help our content rank higher for relevant search queries in the AI Engineering domain.
                    
AI Engineering Academy specializes in:
- Practical AI/ML system design and implementation
- Machine learning engineering workflows
- MLOps and production AI systems
- AI engineering best practices and patterns
- Large Language Model (LLM) fine-tuning and deployment
- AI application architecture and scalability
- Responsible AI engineering and ethics

When analyzing content, identify and emphasize:
1. Primary AI engineering concepts, tools, or methodologies discussed
2. Practical applications and implementation strategies
3. Technical skills being taught
4. Target audience expertise level (beginner, intermediate, advanced)
5. Key problems the content helps solve for AI practitioners
""",
                },
                {"role": "user", "content": content},
                {
                    "role": "user",
                    "content": """Create an SEO-optimized description (160-170 characters) and a longer summary (3-4 sentences) for this content.
                    
For the description:
- Include 2-3 high-value keywords related to AI engineering
- Make it compelling for someone searching for practical AI knowledge
- Ensure it accurately represents the actual content
- Keep it between 160-170 characters (Google search snippet length)
                    
For the summary:
- Highlight the practical AI engineering skills or knowledge gained
- Include 5-7 relevant keywords that someone might search for when looking for this information
- Mention specific tools, frameworks, or methodologies covered (if any)
- Emphasize the problem-solving aspect for AI practitioners
- Structure as 3-4 concise, information-rich sentences

Return your response in this exact YAML format:
description: "Your SEO-optimized description here (160-170 chars)"
summary: "Your longer 3-4 sentence summary here with key AI engineering terms"
keywords: ["keyword1", "keyword2", "keyword3", "keyword4", "keyword5"]
"""
                },
            ],
            max_tokens=4000,
        )
        return response.choices[0].message.content
    except Exception as e:
        console.print(f"[bold red]Error summarizing {path}: {str(e)}[/bold red]")
        return ""


async def generate_sitemap(
    root_dir: str, output_file: str, api_key: Optional[str] = None
) -> None:
    """
    Generate a sitemap from the given root directory.

    Args:
        root_dir (str): The root directory to start traversing from.
        output_file (str): The output file to save the sitemap.
        api_key (Optional[str]): The OpenAI API key. If not provided, it will be read from the OPENAI_API_KEY environment variable.
    """
    client = AsyncOpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY"))

    # Load existing sitemap if it exists
    existing_sitemap: Dict[str, Dict[str, str]] = {}
    if os.path.exists(output_file):
        with open(output_file, "r", encoding="utf-8") as sitemap_file:
            existing_sitemap = yaml.safe_load(sitemap_file) or {}

    sitemap_data: Dict[str, Dict[str, str]] = {}

    async def process_file(
        path: str, content: str, content_hash: str
    ) -> Tuple[str, Dict[str, str]]:
        if (
            path in existing_sitemap
            and existing_sitemap[path].get("hash") == content_hash
        ):
            return path, existing_sitemap[path]
        summary = await summarize_content(client, path, content)
        return path, {"summary": summary, "hash": content_hash}

    with Progress() as progress:
        task = progress.add_task(
            "[green]Processing files...", total=len(list(traverse_docs(root_dir)))
        )
        tasks = [
            process_file(path, content, content_hash)
            for path, content, content_hash in traverse_docs(root_dir)
        ]
        results = await asyncio.gather(*tasks)
        for _ in results:
            progress.update(task, advance=1)

    sitemap_data = dict(results)

    with open(output_file, "w", encoding="utf-8") as sitemap_file:
        yaml.dump(sitemap_data, sitemap_file, default_flow_style=False)

    console.print(
        f"[bold green]Sitemap has been generated and saved to {output_file}[/bold green]"
    )


app = typer.Typer()


@app.command()
def main(
    root_dir: str = typer.Option("docs", help="Root directory to traverse"),
    output_file: str = typer.Option("sitemap.yaml", help="Output file for the sitemap"),
    api_key: Optional[str] = typer.Option(None, help="OpenAI API key"),
):
    """
    Generate a sitemap from the given root directory.
    """
    asyncio.run(generate_sitemap(root_dir, output_file, api_key))


if __name__ == "__main__":
    app()